import psycopg2
from django.db.utils import load_backend, DatabaseError
from .introspection import TenantSchemaDatabaseIntrospection

psql_backend = load_backend('django.db.backends.postgresql')


class DatabaseWrapper(psql_backend.DatabaseWrapper):
    """
    参考自: `django-tenants <https://django-tenants.readthedocs.io/en/latest/index.html>`_

    拓展了Django的原生PostgreSQL数据库Wrapper.
    根据.schemas的具体值，在获取数据库指针时，自动执行"SET search_path TO"

    例子::

        from django.db import connection
        connection.set_tenant('zmm_dorabot')
        # 在这之后的所有ORM操作将只能接触到public和zmm_dorabot模式
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.is_search_path_set = False
        self.schemas = ('public',)  # The current ORM scope

        
        # Introspection预实现了一系列的SQL语句，使得Django可以检查表、数据库、用户的各种状态。
        # 我们需使用一个经定制的Introspection实现，以将检查局限到当前设置的schema。
        self.introspection = TenantSchemaDatabaseIntrospection(self)

    @property
    def primary_schema(self):
        return self.schemas[0]

    def close(self):
        """
        关闭数据库连接，并标注search_path需要重新设置
        """
        super().close()
        self.is_search_path_set = False

    def rollback(self):
        """
        Django执行rollback操作后会清空search_path。因此，我们需要标注search_path需要重新设置
        """
        super().rollback()
        self.is_search_path_set = False

    def set_schema(self, *schemas, include_public=True):
        """
        设置当前数据库会话的可访问schema，默认也允许访问pubic schema
        """
        self.schemas = schemas
        if include_public and 'public' not in schemas:
            self.schemas += ('public',)

        self.is_search_path_set = False

    def _cursor(self, name=None):
        """
        根据.schemas的具体值，在获取数据库指针时，自动执行"SET search_path TO"
        """
        cursor = super()._cursor(name=name)

        if self.is_search_path_set:
            return cursor

        try:
            search_paths_str = ','.join(self.schemas)
            cursor.execute(f'SET search_path = {search_paths_str}')
        except (DatabaseError, psycopg2.InternalError):
            self.is_search_path_set = False
        else:
            self.is_search_path_set = True

        return cursor
